package com.xjgzinfo.usdp.filter;

import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.util.concurrent.RateLimiter;
import org.springframework.stereotype.Component;
import org.springframework.util.AntPathMatcher;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;

@Component
public class QpsFilter implements Filter {
    private static final Map<String, Double> URL_QPS_MAP = new HashMap<>();
    private static final double DEFAULT_URL_QPS = 2000.0;
    private static final LoadingCache<String, RateLimiter> UPLOAD_LIMIT_CACHE = CacheBuilder.newBuilder()
            .maximumSize(1000)
            .refreshAfterWrite(20, TimeUnit.MINUTES)
            .build(new CacheLoader<String, RateLimiter>() {
                @Override
                public RateLimiter load(String key) {
                    //restful
                    AntPathMatcher antPathMatcher = new AntPathMatcher();
                    Optional<Double> result = URL_QPS_MAP.entrySet()
                            .stream()
                            .filter(it -> antPathMatcher.match(it.getKey(), key))
                            .map(Map.Entry::getValue)
                            .findFirst();
                    if (result.isPresent()) {
                        return RateLimiter.create(result.get());
                    }
                    return RateLimiter.create(DEFAULT_URL_QPS);
                }
            });

    static {
        //减少sso压力，2秒只能上传一张图
        URL_QPS_MAP.put("/api/visual/put-file", 1.0);
        URL_QPS_MAP.put("/api/page-schemas/{id}/upload-page-image/{schemaId}", 1.0);
    }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
        HttpServletRequest servletRequest = (HttpServletRequest) request;
        String uri = servletRequest.getRequestURI();
        try {
            RateLimiter rateLimiter = UPLOAD_LIMIT_CACHE.get(uri);
            if (rateLimiter.tryAcquire()) {
                chain.doFilter(request, response);
            } else {
                HttpServletResponse servletResponse = (HttpServletResponse) response;
                servletResponse.setStatus(400);
                servletResponse.getWriter().flush();
            }
        } catch (ExecutionException e) {

        }
    }
}
